﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Diagnostics;
using System.Threading;
using Drifting.MotionPlanning.Primitives;
using Drifting.Foundation;

namespace Drifting.MotionPlanning.PointSets
{
    public class Bin
    {
        public int key_level;
        public double[] key_min;
        public double[] key_max;
        public double delta;
        
        public int count;
        public List<Bin> sections;
        public List<State> values;

        public ReaderWriterLockSlim valueLock;

        public void AddState(State x)
        {
            if (sections != null)
            {
                int dof_x = x.Alias.Length;
                int key_index = key_level;
                while (key_index >= dof_x)
                    key_index -= dof_x;

                int rasterization = (int)Math.Floor((x.Alias[key_index] - key_min[key_index]) / delta);
                Bin subBin = sections[rasterization];
                subBin.AddState(x);
            }
            else
            {
                valueLock.EnterWriteLock();

                values.Add(x);
                x.Index = values.Count - 1;

                valueLock.ExitWriteLock();
            }
            count++;
        }

        public void DelState(State x)
        {
            if (count == 0)
                return;
            if (sections != null)
            {
                int dof_x = x.Alias.Length;
                int key_index = key_level;
                while (key_index >= dof_x)
                    key_index -= dof_x;

                int rasterization = (int)Math.Floor((x.Alias[key_index] - key_min[key_index]) / delta);
                Bin subBin = sections[rasterization];
                subBin.DelState(x);
                count--;
            }
            else
            {
                valueLock.EnterWriteLock();

                int idx_del = x.Index;
                if (idx_del >= 0)
                {
                    int idx_last = values.Count - 1;
                    if (idx_del < idx_last)
                    {
                        State x_last = values[idx_last];
                        values[idx_del] = x_last;
                        x_last.Index = idx_del;
                    }
                    values.RemoveAt(idx_last);
                    x.Index = -1;

                    count--;
                }

                valueLock.ExitWriteLock();
            }

        }

        public int DelState(double cost_threshold)
        {
            if (values == null)
            {
                int count_new = 0;
                foreach (Bin section in sections)
                {
                    count_new += section.DelState(cost_threshold);
                }
                count = count_new;
            }
            else
            {
                valueLock.EnterWriteLock();


                for (int i = 0; i < values.Count;++i )
                {
                    State x = values[i];
                    if (x != null)
                        if (x.Cost > cost_threshold)
                        {
                            int idx_last = values.Count - 1;
                            State x_last = values[idx_last];
                            values[i] = x_last;
                            x_last.Index = i;
                            x.Index = -1;
                            values.RemoveAt(idx_last);
                        }
                }

                count = values.Count;

                valueLock.ExitWriteLock();
            }
            return count;
        }

        public void TraverseForTheClosest(double[] key, ref State x_nearest, ref double dist2)
        {
            if (values != null)
            {
                valueLock.EnterReadLock();

                int i = 0;
                State x_i = null;

                while(i<values.Count)
                {
                    x_i = values[i];

                    if (x_i != null)
                    {
                        double dist_i = PointSet.Euclidean2(key, x_i.Alias);
                        if (dist_i < dist2)
                        {
                            dist2 = dist_i;
                            x_nearest = x_i;
                        }
                    }
                    else
                    {
                        throw new InvalidProgramException(" this should not happen ");
                    }
                    i++;
                }

                valueLock.ExitReadLock();
            }
            else
            {
                foreach (Bin b_i in sections)
                {
                    if(b_i.count>0)
                        b_i.TraverseForTheClosest(key, ref x_nearest, ref dist2);
                }
            }
        }

        public Bin DeepestNonEmptyBin(double[] key)
        {
            if (count == 0)
                return null;
            if (sections != null)
            {
                int dof_x = key.Length;
                int key_index = key_level;
                while (key_index >= dof_x)
                    key_index -= dof_x;

                int rasterization = (int)Math.Floor((key[key_index] - key_min[key_index]) / delta);
                Bin subBin = sections[rasterization].DeepestNonEmptyBin(key);

                if (subBin == null)
                    return this;
                else
                    return subBin;
            }
            else
                return this;
        }

        public void ReportVerticesSnapShot(ref List<State> snapshot, ref int idx, ref double sumX2, ref double sumX, ref int nNonEmpty)
        {
            if (values != null)
            {
                valueLock.EnterReadLock();
                int i = 0;
                while (i < count)
                {
                    if (snapshot.Count <= idx)
                        snapshot.Add(null);
                    if (i >= values.Count) break;
                    snapshot[idx] = values[i];
                    idx++;
                    i++;
                }
                valueLock.ExitReadLock();

                if (count > 0)
                {
                    sumX2 += count * count;
                    sumX += count;
                    nNonEmpty++;
                }
            }
            else
            {
                foreach (Bin b_i in sections)
                {
                    if (b_i.count > 0)
                        b_i.ReportVerticesSnapShot(ref snapshot, ref idx, ref sumX2, ref sumX, ref nNonEmpty);
                    
                }
            }

        }

        public void ReportVerticesSnapShot(ref List<State> snapshot, ref int idx)
        {
            if (values != null)
            {
                valueLock.EnterReadLock();
                int i = 0;
                while (i < count)
                {
                    if (snapshot.Count <= idx)
                        snapshot.Add(null);
                    if (i >= values.Count) break;
                    snapshot[idx] = values[i];
                    idx++;
                    i++;
                }
                valueLock.ExitReadLock();
            }
            else
            {
                foreach (Bin b_i in sections)
                {
                    if (b_i.count > 0)
                        b_i.ReportVerticesSnapShot(ref snapshot, ref idx);

                }
            }

        }

        public void RangeExclusiveTraversal(double[] key, double r, ref List<State> neighbors)
        {
            if (count <= 0) return;

            if (values != null)
            {
                valueLock.EnterReadLock();
                double r2 = r*r;

                int i = 0;
                while(i<values.Count)
                {
                    State x = values[i];
                    if (x != null)
                    {
                        double dist2 = PointSet.Euclidean2(key, x.Alias);
                        if (dist2 <= r2)
                            neighbors.Add(x);
                    }
                    else
                    {
                        throw new InvalidProgramException("should not happen");
                    }
                    i++;
                }
                valueLock.ExitReadLock();
            }
            else
            {
                int dof_x = key.Length;
                int key_index = key_level;
                while (key_index >= dof_x)
                    key_index -= dof_x;

                double b_i_min = key_min[key_index];
                double b_i_max = b_i_min + delta;
                double range_min = key[key_index] - r;
                double range_max = key[key_index] + r;

                for (int i = 0; i<sections.Count; ++i)
                {
                    if (b_i_max > range_min && b_i_min < range_max)
                        sections[i].RangeExclusiveTraversal(key, r, ref neighbors);

                    b_i_min += delta;
                    b_i_max += delta;
                }
            }
        }
    }

    public class BinQuery : PointSet
    {
        double[] alias_min;
        double[] alias_max;
        int alias_dof;
        int levels;
        int[] branches;

        Bin root;
        int n_value_bins;

        override public int Count
        { get { return root.count; } }

        public BinQuery(double[] alias_min, double[] alias_max, int alias_dof, int[] branches)
        {
            this.alias_max = alias_max;
            this.alias_min = alias_min;
            this.alias_dof = alias_dof;
            this.branches = branches;
            this.levels = branches.Length;
            root = new Bin();
            root.key_level = 0;
            root.key_min = alias_min;
            root.key_max = alias_max;
            root.count = 0;
            //add information to root

            n_value_bins = 1;
            foreach (int b in branches)
                n_value_bins *= b;

            ConstructBinHierarchy(root);
        }

        public int NumberValueBins
        { get { return n_value_bins; } }
        
        public void ConstructBinHierarchy(Bin bin)
        {
            if (bin.key_level < levels)
            {
                bin.sections = new List<Bin>();

                int cur_index = bin.key_level;
                while (cur_index >= alias_dof)
                    cur_index -= alias_dof;

                int n_branch = branches[cur_index];
                double range_level = bin.key_max[cur_index] - bin.key_min[cur_index];
                double delta_level = range_level / n_branch;
                bin.delta = delta_level;

                for (int i = 0; i < n_branch; ++i)
                {
                    Bin b_i = new Bin();
                    bin.sections.Add(b_i);

                    b_i.key_level = bin.key_level + 1;
                    b_i.key_min = new double[alias_dof];
                    b_i.key_max = new double[alias_dof];
                    MathAssistant.vector_set(bin.key_min, ref b_i.key_min);
                    MathAssistant.vector_set(bin.key_max, ref b_i.key_max);
                    b_i.key_min[cur_index] = bin.key_min[cur_index] + i * delta_level;
                    b_i.key_max[cur_index] = b_i.key_min[cur_index] + delta_level;

                    ConstructBinHierarchy(b_i);
                }
            }
            else
            {
                bin.values = new List<State>(16);
                bin.sections = null;
                bin.valueLock = new ReaderWriterLockSlim();
            }
        }

        override public bool inSpace(double[] key)
        {
            if (key.Length != alias_dof)
                throw new ArgumentException("key.dimension incorrect");
            for (int i = 0; i < alias_dof; ++i)
            {
                if (key[i] < alias_min[i] || key[i] > alias_max[i])
                    return false;
            }
            return true;
        }

        override public bool AddState(State x)
        {
            if (root == null)
                throw new InvalidOperationException("Forget to call ConstructBinHierarchy?");

            if (inSpace(x.Alias) == false)
                return false;

            root.AddState(x);

            return true;
        }

        override public bool DelState(State x)
        {
            if (x == null)
                return false;

            if (inSpace(x.Alias) == false)
                return false;

            root.DelState(x);

            return true;
        }

        public void DelState(double cost_threshold)
        {
            root.DelState(cost_threshold);
        }

        override public void Neighbor_SingleQuery(double[] key, ref State x_nearest, ref double dist2)
        {
            if (inSpace(key) == false)
            {
                // if key is not the alias, an exception will be thrown
                dist2 = double.MaxValue;
                x_nearest = null;
                return;
            }

            Bin deepestNonEmptyBin = root.DeepestNonEmptyBin(key);
            LinearScreenTheNearest(deepestNonEmptyBin, key, ref x_nearest, ref dist2);
        }

        override public void Neighbor_RangeQuery(double[] key, double r, ref List<State> neighbors)
        {
            if (inSpace(key) == false)
            {
                neighbors.Clear();
                return;
            }
            
            if (neighbors == null)
                neighbors = new List<State>(64);
            neighbors.Clear();

            root.RangeExclusiveTraversal(key, r, ref neighbors);
        }

        override public void Neighbor_KQuery(double[] key, int K, ref List<State> neighbors)
        {
            double r = 0.1;
            neighbors.Clear();

            while (neighbors.Count < K)
            {
                Neighbor_RangeQuery(key, r, ref neighbors);

                if (neighbors.Count >= root.count)
                    break;

                r *= 2;
            }

            BinaryMinHeap minHeap = new BinaryMinHeap(neighbors.Count);
            foreach (State x in neighbors)
                minHeap.Push(PointSet.Euclidean2(key, x.Q), x);

            neighbors.Clear();
            for (int i = 0; i < K; ++i)
            {
                double dist = 0;
                object value = null;
                minHeap.Pop(ref dist, ref value);
                if (value != null)
                    neighbors.Add((State)value);
            }
        }

        protected void LinearScreenTheNearest(Bin bin, double[] key, ref State x_nearest, ref double dist2)
        {
            dist2 = double.MaxValue;
            bin.TraverseForTheClosest(key, ref x_nearest, ref dist2);
        }

        override public void ReportVerticesSnapshot(ref List<State> snapshot)
        {
            if (snapshot == null )
            {
                snapshot = new List<State>(root.count);
                for (int i = 0; i < root.count; ++i)
                    snapshot.Add(null);
            }
            else if(snapshot.Count < root.count)
            {
                while (snapshot.Count < root.count)
                    snapshot.Add(null);
            }
            int idx = 0;
            root.ReportVerticesSnapShot(ref snapshot, ref idx);
        }

        public void ReportVerticesSnapshot(ref List<State> snapshot, ref double mean, ref double var, ref int nNonEmpty)
        {
            if (snapshot == null)
            {
                snapshot = new List<State>(root.count);
                for (int i = 0; i < root.count; ++i)
                    snapshot.Add(null);
            }
            else if (snapshot.Count < root.count)
            {
                while (snapshot.Count < root.count)
                    snapshot.Add(null);
            }
            double sumX2 = 0;
            double sumX = 0;
            int idx = 0;
            root.ReportVerticesSnapShot(ref snapshot, ref idx, ref sumX2, ref sumX, ref nNonEmpty);

            if (nNonEmpty != 0)
            {
                mean = sumX / nNonEmpty;
                var = sumX2 / nNonEmpty - mean * mean;
            }
        }
    }
}
